import os

import dotenv
import requests
import json

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import tool, Tool, BaseTool

from internal.lib.helper import add_attribute

dotenv.load_dotenv()

class QianFanArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")

@tool
def baidu_qianfan_img_generate(query: str) -> str:
    """这是百度千帆的ai生成图片的工具。输出是图片的路径"""
    url = "https://qianfan.baidubce.com/v2/images/generations"
    payload = json.dumps({
        "prompt": query,
        "model": "irag-1.0"
    }, ensure_ascii=False)
    gaode_api_key = os.getenv("QIANFAN_API_KEY")
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {gaode_api_key}'
    }
    response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))
    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")

@add_attribute("args_schema", QianFanArgsSchema)
def qianfan_generator(**kwargs) -> BaseTool:
    """创建千帆文生图工具"""
    return Tool(
        name="QianfanImgGenerator",
        func=baidu_qianfan_img_generate,
        description="使用QIAN FAN irag-1.0生成图片",
        args_schema=QianFanArgsSchema
    )

